// Copyright (c) Lawrence Livermore National Security, LLC and other VisIt
// Project developers.  See the top-level LICENSE file for dates and other
// details.  No copyright assignment is required to contribute to VisIt.

// ************************************************************************* //
//                            avtXDATFileFormat.C                           //
// ************************************************************************* //

#include <avtXDATFileFormat.h>

#include <float.h>
#include <string>

#include <vtkFloatArray.h>
#include <vtkDoubleArray.h>
#include <vtkPolyData.h>
#include <vtkCellArray.h>

#include <avtDatabaseMetaData.h>
#include <avtMTSDFileFormatInterface.h>
#include <AtomicProperties.h>

#include <Expression.h>

#include <InvalidVariableException.h>
#include <InvalidFilesException.h>
#include <sstream>

#include <vtkTriangulationTables.h>

using     std::string;

// ****************************************************************************
//  Method: avtXDAT constructor
//
//  Programmer: Jeremy Meredith
//  Creation:   January  8, 2008
//
//  Modifications:
//
// ****************************************************************************

avtXDATFileFormat::avtXDATFileFormat(const char *fn)
    : avtMTSDFileFormat(&fn, 1)
{
    filename = fn;
    OpenFileAtBeginning();

    metadata_read = false;
    full_header = false;

    natoms = 0;

    unitCell[0][0] = 1;    unitCell[0][1] = 0;    unitCell[0][2] = 0;
    unitCell[1][0] = 0;    unitCell[1][1] = 1;    unitCell[1][2] = 0;
    unitCell[2][0] = 0;    unitCell[2][1] = 0;    unitCell[2][2] = 1;
}


// ****************************************************************************
//  Method: avtXDATFileFormat::FreeUpResources
//
//  Purpose:
//      When VisIt is done focusing on a particular timestep, it asks that
//      timestep to free up any resources (memory, file descriptors) that
//      it has associated with it.  This method is the mechanism for doing
//      that.
//
//  Programmer: Jeremy Meredith
//  Creation:   January  8, 2008
//
//  Modifications:
//
// ****************************************************************************

void
avtXDATFileFormat::FreeUpResources(void)
{
    x.clear();
    y.clear();
    z.clear();
    species.clear();
    species_counts.clear();
    element_map.clear();
    metadata_read = false;
    full_header = false;
}


// ****************************************************************************
//  Method:  avtXDATFileFormat::GetNTimesteps
//
//  Purpose:
//    return the number of timesteps
//
//  Arguments:
//    none
//
//  Programmer:  Jeremy Meredith
//  Creation:    March  4, 2016
//
//  Modifications:
//
// ****************************************************************************
int
avtXDATFileFormat::GetNTimesteps(void)
{
    ReadMetaData();
    return file_positions.size();
}

// ****************************************************************************
//  Method: avtXDATFileFormat::PopulateDatabaseMetaData
//
//  Purpose:
//      This database meta-data object is like a table of contents for the
//      file.  By populating it, you are telling the rest of VisIt what
//      information it can request from you.
//
//  Programmer: Jeremy Meredith
//  Creation:   January  8, 2008
//
//  Modifications:
//    Jeremy Meredith, Tue Feb 12 14:09:24 EST 2008
//    Support element types as an enumerated scalar.
//
//    Mark C. Miller, Mon Apr 14 15:41:21 PDT 2008
//    Changed interface to enum scalars
//
//    Jeremy Meredith, Thu Oct 18 10:58:20 EDT 2012
//    Changed enum names to be origin H=1.
//
// ****************************************************************************

void
avtXDATFileFormat::PopulateDatabaseMetaData(avtDatabaseMetaData *md, int)
{
    ReadMetaData();

    avtMeshMetaData *mmd = new avtMeshMetaData("mesh", 1, 0,0,0,
                                               3, 0,
                                               AVT_POINT_MESH);
    mmd->nodesAreCritical = true;
    for (int i=0; i<9; i++)
    {
        mmd->unitCellVectors[i] = unitCell[i/3][i%3];
    }
    md->Add(mmd);

    avtMeshMetaData *mmd_bbox = new avtMeshMetaData("unitCell", 1, 0,0,0,
                                                    3, 1,
                                                    AVT_UNSTRUCTURED_MESH);
    for (int i=0; i<9; i++)
    {
        mmd_bbox->unitCellVectors[i] = unitCell[i/3][i%3];
    }
    md->Add(mmd_bbox);

    AddScalarVarToMetaData(md, "species", "mesh", AVT_NODECENT);
    if (element_map.size() > 0)
    {
        avtScalarMetaData *el_smd =
            new avtScalarMetaData("element", "mesh", AVT_NODECENT);
        el_smd->SetEnumerationType(avtScalarMetaData::ByValue);
        for (size_t i=0; i<element_map.size(); i++)
            el_smd->AddEnumNameValue(element_names[element_map[i]],element_map[i]);
        md->Add(el_smd);
    }
    if (cx.size() > 0)
        AddScalarVarToMetaData(md, "cx", "mesh", AVT_NODECENT);
    if (cy.size() > 0)
        AddScalarVarToMetaData(md, "cy", "mesh", AVT_NODECENT);
    if (cz.size() > 0)
        AddScalarVarToMetaData(md, "cz", "mesh", AVT_NODECENT);
}


// ****************************************************************************
//  Method: avtXDATFileFormat::GetMesh
//
//  Purpose:
//      Gets the mesh associated with this file.  The mesh is returned as a
//      derived type of vtkDataSet (ie vtkRectilinearGrid, vtkStructuredGrid,
//      vtkUnstructuredGrid, etc).
//
//  Arguments:
//      meshname    The name of the mesh of interest.  This can be ignored if
//                  there is only one mesh.
//
//  Programmer: Jeremy Meredith
//  Creation:   January  8, 2008
//
//  Modifications:
//
// ****************************************************************************

vtkDataSet *
avtXDATFileFormat::GetMesh(int ts, const char *name)
{
    ReadTimestep(ts);

    string meshname(name);
    if (meshname == "unitCell")
    {
        vtkPolyData *pd  = vtkPolyData::New();
        vtkPoints   *pts = vtkPoints::New();

        pts->SetNumberOfPoints(8);
        pd->SetPoints(pts);
        pts->Delete();
        for (int j = 0 ; j < 8 ; j++)
        {
            float x=0,y=0,z=0;
            for (int axis=0; axis<3; axis++)
            {
                if (j & (1<<axis))
                {
                    x += unitCell[axis][0];
                    y += unitCell[axis][1];
                    z += unitCell[axis][2];
                }
            }
            pts->SetPoint(j, x,y,z);
        }
 
        vtkCellArray *lines = vtkCellArray::New();
        pd->SetLines(lines);
        lines->Delete();
        for (int k = 0 ; k < 12 ; k++)
        {
            lines->InsertNextCell(2);
            lines->InsertCellPoint(voxVerticesFromEdges[k][0]);
            lines->InsertCellPoint(voxVerticesFromEdges[k][1]);
        }

        return pd;
    }
    else if (meshname == "mesh")
    {
        vtkPolyData *pd  = vtkPolyData::New();
        vtkPoints   *pts = vtkPoints::New();

        pts->SetNumberOfPoints(natoms);
        pd->SetPoints(pts);
        pts->Delete();
        for (int j = 0 ; j < natoms ; j++)
        {
            pts->SetPoint(j, x[j], y[j], z[j]);
        }
 
        vtkCellArray *verts = vtkCellArray::New();
        pd->SetVerts(verts);
        verts->Delete();
        for (int k = 0 ; k < natoms ; k++)
        {
            verts->InsertNextCell(1);
            verts->InsertCellPoint(k);
        }

        return pd;
    }

    return NULL;
}


// ****************************************************************************
//  Method: avtXDATFileFormat::GetVar
//
//  Purpose:
//      Gets a scalar variable associated with this file.  Although VTK has
//      support for many different types, the best bet is vtkFloatArray, since
//      that is supported everywhere through VisIt.
//
//  Arguments:
//      varname    The name of the variable requested.
//
//  Programmer: Jeremy Meredith
//  Creation:   January  8, 2008
//
//  Modifications:
//
// ****************************************************************************

vtkDataArray *
avtXDATFileFormat::GetVar(int ts, const char *varname)
{
    ReadTimestep(ts);

    if (string(varname) == "species")
    {
        vtkFloatArray *scalars = vtkFloatArray::New();
        scalars->SetNumberOfTuples(natoms);
        float *ptr = (float *) scalars->GetVoidPointer(0);
        for (int i=0; i<natoms; i++)
        {
            ptr[i] = species[i];
        }
        return scalars;
    }

    if (string(varname) == "element")
    {
        vtkFloatArray *scalars = vtkFloatArray::New();
        scalars->SetNumberOfTuples(natoms);
        float *ptr = (float *) scalars->GetVoidPointer(0);
        for (int i=0; i<natoms; i++)
        {
            ptr[i] = element_map[species[i]];
        }
        return scalars;
    }

    if (string(varname) == "cx")
    {
        vtkFloatArray *scalars = vtkFloatArray::New();
        scalars->SetNumberOfTuples(natoms);
        float *ptr = (float *) scalars->GetVoidPointer(0);
        for (int i=0; i<natoms; i++)
        {
            ptr[i] = cx[i];
        }
        return scalars;
    }

    if (string(varname) == "cy")
    {
        vtkFloatArray *scalars = vtkFloatArray::New();
        scalars->SetNumberOfTuples(natoms);
        float *ptr = (float *) scalars->GetVoidPointer(0);
        for (int i=0; i<natoms; i++)
        {
            ptr[i] = cy[i];
        }
        return scalars;
    }

    if (string(varname) == "cz")
    {
        vtkFloatArray *scalars = vtkFloatArray::New();
        scalars->SetNumberOfTuples(natoms);
        float *ptr = (float *) scalars->GetVoidPointer(0);
        for (int i=0; i<natoms; i++)
        {
            ptr[i] = cz[i];
        }
        return scalars;
    }

    return NULL;
}


// ****************************************************************************
//  Method: avtXDATFileFormat::GetVectorVar
//
//  Purpose:
//      Gets a vector variable associated with this file.  Although VTK has
//      support for many different types, the best bet is vtkFloatArray, since
//      that is supported everywhere through VisIt.
//
//  Arguments:
//      varname    The name of the variable requested.
//
//  Programmer: Jeremy Meredith
//  Creation:   January  8, 2008
//
//  Modifications:
//
// ****************************************************************************

vtkDataArray *
avtXDATFileFormat::GetVectorVar(int ts, const char *varname)
{
    return NULL;
}


// ****************************************************************************
//  Method:  avtXDATFileFormat::ReadMetaData
//
//  Purpose:
//    Read the metadata and atoms.
//
//  Arguments:
//    none
//
//  Programmer:  Jeremy Meredith
//  Creation:    January  8, 2008
//
//  Modifications:
//   Jeremy Meredith, Tue Dec 29 13:41:16 EST 2009
//   Added some error checks.
//
//    Jeremy Meredith, Mon Jan 24 17:03:27 EST 2011
//    In newer VASP flavors, there is an optional line with atomic symbols
//    above the line which lists the counts of atoms for each species
//    type.  This line takes precedence over any "types=" line in the header.
//
// ****************************************************************************
void
avtXDATFileFormat::ReadHeader( bool read_full_header )
{
    char line[132];

    if( read_full_header )
    {
      // this line is the comment line.
      // At the moment, species look like they are on a different 
      // line, so I am not attempting to parse this line for an "="
      // symbol which describes species.  If we want to add it back
      // in, then pick up the code from the POSCAR reader.
      in.getline(line, 132);
      
      // Read the scale and lattice vectors
      in >> scale;
      in >> lat[0][0] >> lat[0][1] >> lat[0][2];
      in >> lat[1][0] >> lat[1][1] >> lat[1][2];
      in >> lat[2][0] >> lat[2][1] >> lat[2][2];
      
      // error check
      if (scale == 0)
        EXCEPTION2(InvalidFilesException, filename.c_str(), "Scale was not "
                   "a nonzero real number; does not match XDAT format.");

      for (int i=0; i<3; i++)
        for (int j=0; j<3; j++)
          unitCell[i][j] = scale*lat[i][j];

      in.getline(line, 132); // skip rest of the last lattice line
      
      // get atom counts, and optionally, element types
      in.getline(line, 132);
      string atomtypeline(line);
      string atomcountline(line);
      
      std::istringstream type_in(atomtypeline);
      string tmp_element;

      if ((type_in >> tmp_element) &&
          ElementNameToAtomicNumber(tmp_element.c_str()) > 0)
      {
        // We've got an element types line to parse.
        // This overrides any earlier "types =" from above:
        element_map.clear();
        // Push the one we already read:
        element_map.push_back(ElementNameToAtomicNumber(tmp_element.c_str()));
        // Read the rest:
        while (type_in >> tmp_element)
        {
          element_map.push_back(ElementNameToAtomicNumber(tmp_element.c_str()));
        }
        
        // We need to read the next line for the atom counts: we set it up
        // to use this past line in the event we didn't have a species line:
        in.getline(line, 132);
        atomcountline = line;
      }

      natoms = 0;
      int tmp_count;
      species_counts.clear();
      std::istringstream count_in(atomcountline);
      while (count_in >> tmp_count)
      {
        species_counts.push_back(tmp_count);
        natoms += tmp_count;
      }

      // error check
      if (natoms == 0)
        EXCEPTION2(InvalidFilesException, filename.c_str(),
                   "Could not parse atom counts; does not match XDAT format.");

      // If we tried to create an element map, but it did not have the
      // same number of values as the species counts, it is invalid.
      if (element_map.size() > 0 && element_map.size() != species_counts.size())
      {
        element_map.clear();
      }
    }
    
    // next line is either Selective dynamics or Cartesian/Direct
    selective_dynamics = false;
    in.getline(line, 132);
    if (line[0] == 's' || line[0] == 'S')
    {
        selective_dynamics = true;
        // if it was Selective dynamics, then Cartesian/Direct is the next one
        in.getline(line, 132);
    }

    cartesian = false;
    if (line[0] == 'c' || line[0] == 'C' || line[0] == 'k' || line[0] == 'K')
    {
        cartesian = true;
    }
    // if it's not c/k, then it's "Direct"
}


// ****************************************************************************
//  Method:  avtXDATFileFormat::ReadMetaData
//
//  Purpose:
//    Read the metadata and atoms.
//
//  Arguments:
//    none
//
//  Programmer:  Jeremy Meredith
//  Creation:    January  8, 2008
//
//  Modifications:
//   Jeremy Meredith, Tue Dec 29 13:41:16 EST 2009
//   Added some error checks.
//
//    Jeremy Meredith, Mon Jan 24 17:03:27 EST 2011
//    In newer VASP flavors, there is an optional line with atomic symbols
//    above the line which lists the counts of atoms for each species
//    type.  This line takes precedence over any "types=" line in the header.
//
// ****************************************************************************
void
avtXDATFileFormat::ReadMetaData()
{
    if (metadata_read)
        return;

    metadata_read = true;

    // New format has a full header; the title, scale, unit cell,
    // species, and number of each species and other information which
    // is at each time step
    char line[132];
    
    // This line is the title line.
    // At the moment, species look like they are on a different 
    // line, so I am not attempting to parse this line for an "="
    // symbol which describes species.  If we want to add it back
    // in, then pick up the code from the POSCAR reader.
    in.getline(line, 132);
    string titleline(line);

    // Reset to the begininng of the file to read the full header.
    in.seekg(0);
    
    ReadHeader(1);

    // Read all of the atoms for the first time step.
    for (int i=0; i<natoms; i++)
      in.getline(line,132);

    // Get the next line and check for a full header.
    in.getline(line, 132);

    if( string(line) == titleline )
      full_header = true;

    // Reset to the begininng of the file to read each time step
    in.seekg(0);
    
    // count the timesteps!
    while( in.peek() != EOF )
    {
        // If a full header point to that start point.
        if( full_header )
          file_positions.push_back(in.tellg());

        // Read the full header for the first time step.
        ReadHeader(file_positions.size() == 0 || full_header );
        
        // If a partial header point to the first atom after the header.
        if( !full_header )
          file_positions.push_back(in.tellg());

        // Read all of the atoms.
        for (int i=0; i<natoms; i++)
            in.getline(line,132);
    }
}

void
avtXDATFileFormat::ReadTimestep(int timestep)
{
    ReadMetaData();

    OpenFileAtBeginning();
    in.seekg(file_positions[timestep]);
    if( full_header )
      ReadHeader( full_header );
      
    char line[132];

    // read the atoms
    x.resize(natoms);
    y.resize(natoms);
    z.resize(natoms);
    species.resize(natoms);
    if (selective_dynamics)
    {
        cx.resize(natoms);
        cy.resize(natoms);
        cz.resize(natoms);
    }

    int species_index = 0;
    int species_count = 0;
    for (int i=0; i<natoms; i++)
    {
        in.getline(line,132);
        string atomline(line);
        std::istringstream atom_in(atomline);

        double tx, ty, tz;
        atom_in >> tx >> ty >> tz;
        if (cartesian)
        {
            x[i] = tx * scale;
            y[i] = ty * scale;
            z[i] = tz * scale;
        }
        else
        {
            x[i] = tx*unitCell[0][0] + ty*unitCell[1][0] + tz*unitCell[2][0];
            y[i] = tx*unitCell[0][1] + ty*unitCell[1][1] + tz*unitCell[2][1];
            z[i] = tx*unitCell[0][2] + ty*unitCell[1][2] + tz*unitCell[2][2];
        }
        if (selective_dynamics)
        {
            char tmp;
            atom_in >> tmp;
            cx[i] = (tmp=='t' || tmp=='T');
            atom_in >> tmp;
            cy[i] = (tmp=='t' || tmp=='T');
            atom_in >> tmp;
            cz[i] = (tmp=='t' || tmp=='T');
        }
        species[i] = species_index;
        species_count++;
        if (species_count >= species_counts[species_index])
        {
            species_index++;
            species_count = 0;
        }
    }

    in.close();
}

// ****************************************************************************
//  Method:  avtXDATFileFormat::Identify
//
//  Purpose:
//    Return true if the file given is a XDAT VASP file.
//    So far, only check based on the filename.
//
//  Arguments:
//    filename   the filename
//
//  Programmer:  Jeremy Meredith
//  Creation:    January  8, 2008
//
//  Modifications:
//
// ****************************************************************************
bool
avtXDATFileFormat::Identify(const std::string &filename)
{
    // strip off all leading paths
    int pos = filename.length()-1;
    while (pos>=0 && filename[pos]!='/' && filename[pos]!='\\')
        pos--;

    std::string fn;
    if (pos >= 0)
        fn = filename.substr(pos+1);
    else
        fn = filename;

    // uppercase it
    for (size_t i=0; i<fn.size(); i++)
    {
        if (fn[i]>='a' && fn[i]<='z')
            fn[i] = fn[i] + ('A'-'a');
    }

    for (size_t i=0; i<=fn.length()-3; i++)
    {
        if (fn.substr(i,4) == "XDAT")
            return true;
    }

    return false;
}

// ****************************************************************************
//  Method:  avtXDATFileFormat::CreateInterface
//
//  Purpose:
//    Create a file format interface from this reader.
//
//  Programmer:  Jeremy Meredith
//  Creation:    January  8, 2008
//
//  Modifications:
//
// ****************************************************************************
avtFileFormatInterface *
avtXDATFileFormat::CreateInterface(const char *const *list,
                                         int nList, int nBlock)
{
    int nTimestepGroups = nList / nBlock;
    avtMTSDFileFormat ***ffl = new avtMTSDFileFormat**[nTimestepGroups];
    for (int i = 0 ; i < nTimestepGroups ; i++)
    {
        ffl[i] = new avtMTSDFileFormat*[nBlock];
        for (int j = 0 ; j < nBlock ; j++)
        {
            ffl[i][j] = new avtXDATFileFormat(list[i*nBlock+j]);
        }
    }
    return new avtMTSDFileFormatInterface(ffl, nTimestepGroups, nBlock);
}


// ****************************************************************************
//  Method:  avtXDATFileFormat::OpenFileAtBeginning
//
//  Purpose:
//    Opens the file, or else seeks to the beginning.
//
//  Arguments:
//    none
//
//  Programmer:  Jeremy Meredith
//  Creation:    January  8, 2008
//
// ****************************************************************************
void
avtXDATFileFormat::OpenFileAtBeginning()
{
    if (!in.is_open())
    {
        in.open(filename.c_str());
        if (!in)
        {
            EXCEPTION1(InvalidFilesException, filename.c_str());
        }
    }
    else
    {
        in.clear();
        in.seekg(0, ios::beg);
    }
}
